import requests, json, time, sys
from django.conf import settings
from mysqldb.models import AllinpayRobotQiniuLog, AllinpayRobotConfig
from iClouds.apiheader import has_chocode
from iClouds.SDKqiniu import qiMainThead
from gradio_client import Client, handle_file, file
import wave
import soundfile as sf
import numpy as np
import random, os


HOST = "http://120.26.13.46" if settings.DEBUG else "http://172.22.80.140"


class soVitsMainPage:
    def __init__(self):
        pass

    def tts(self, idx):
        """tts推理"""

        # 查询任务
        task = self.hasTasksInfo(idx)
        if not task:
            return {"rescode": 400, "resmsg": "查询任务属性失败"}
        audio = self.hasAudioAttr(task.model)
        if not audio:
            return {"rescode": 400, "resmsg": "查询声音属性失败"}

        # 格式化字符串
        txt = self.doFormatString(task.chat)
        # 计算文本分块
        idx, chunk, local = [], 24, []
        # 按符号分割
        for c in txt.split("。"):
            if len(c) > 1:
                idx.append(f"{c}")

        # 开启TTS
        isTTS = self.Enable_TTS(True, audio.gptmodels, audio.sovitsmodels)
        if not isTTS:
            return {"rescode": 400, "resmsg": "开启TTS服务失败"}

        # 开始循环TTS
        stts = True
        try:
            for g in range(len(idx)):
                print(f"\n{idx[g]}")
                tts = self.soVitsTtssPage(
                    audio.sound, audio.refer, audio.language, idx[g], task.language
                )
                if not tts:
                    raise AssertionError(f"TTS失败 {idx[g]}")
                # 加入列表
                local.append(tts)
        except BaseException as err:
            print("AI服务器TTS失败")
            stts = False
        finally:
            isTTS = self.Enable_TTS(False, audio.gptmodels, audio.sovitsmodels)
        if not stts:
            return {"rescode": 400, "resmsg": "TTS失败"}

        # 合并文件
        if not len(local) == len(idx):
            return {"rescode": 400, "resmsg": "分片编组缺失"}
        gap = self.emptywav()
        arrays = []
        for n in local:
            arrays.append(gap)
            arrays.append(n)
        new_arry = np.concatenate([self.hasarrays(n) for n in arrays])
        out = self.saveOutAudio(new_arry)
        print("分片合并完成")

        # # 提交数据
        # if not settings.DEBUG:
        #     if task.mode in ["Audio", "Voice"]:
        #         task.description = "完成"
        #         task.lasttime = timezone.now()
        #     task.audio = "https://lease.3vppt.com/static" + out.split("statics")[1]
        #     task.status = 1
        #     ttask.save()

        # return {
        #     "rescode": 200,
        #     "resmsg": "https://lease.3vppt.com/static" + out.split("statics")[1],
        # }

        # 上传七牛云
        r = qiMainThead().doUpload(out)
        return {"rescode": 200, "resmsg": f'https://oss.3vppt.com/{r.get("key")}'}

    def Enable_TTS(self, staus, ckpt, pth, gpu="0"):
        client = Client(f"{HOST}:9874")
        result = client.predict(
            staus,  # bool in 'Open TTS inference WEBUI' Checkbox component
            "GPT_SoVITS/pretrained_models/chinese-roberta-wwm-ext-large",
            "GPT_SoVITS/pretrained_models/chinese-hubert-base",
            gpu,  # str in 'GPU number, can only input ONE integer' Textbox component
            ckpt,  # str in '*GPT models list' Dropdown component
            pth,  # str in '*SoVITS models list' Dropdown component
            fn_index=23,
        )
        return (
            True
            if result
            in [
                "TTS inference is enabled",
                "TTS inference process closed",
                "TTS inference process is opened",
            ]
            else False
        )

    def soVitsTtssPage(self, url, refer, language, text, target_language):
        """tts推理"""
        i = 0
        while i < 60:
            try:
                client = Client(f"{HOST}:9872/")
                result = client.predict(
                    url,  # 参考音频url
                    refer,  # 参考音频内容
                    language,  # 参考音频语言
                    f"{text}",  # 要转换的文本
                    target_language,  # 要转换的语言
                    "Cut per 50 characters",  # 切分方式
                    5,
                    1,
                    1,
                    True,
                    fn_index=3,
                )
            except BaseException as err:
                i += 1
            else:
                return result
            time.sleep(1)
        print("连接服务器失败")
        return False

    @staticmethod
    def saveOutAudio(arrays):
        """保存文件"""
        # outfile = (
        #     settings.BASE_DIR.parent
        #     / f'statics/audio/{time.strftime("%Y-%m-%d", time.localtime())}/{has_chocode(16,1).upper()}'
        # )
        # print(outfile)
        # if not os.path.exists(outfile):
        #     os.makedirs(outfile)
        # outfile = os.path.join(outfile, "output.wav")
        outfile = "output.wav"
        with wave.open(outfile, "wb") as f:
            f.setparams((1, 2, 32000, 0, "NONE", "not compressed"))
            f.writeframes(arrays.tobytes())
        return outfile

    @staticmethod
    def emptywav(sample_rate=32000, duration=0.45):
        """创建空白声音"""
        files = "empty.wav"
        num_samples = int(sample_rate * duration)
        array = np.zeros(num_samples, dtype=np.int16)
        sf.write(files, array, sample_rate)
        return files

    @staticmethod
    def hasarrays(url):
        with wave.open(url, "rb") as f:
            params = f.getparams()
            frames = f.readframes(params.nframes)
            array = np.frombuffer(frames, dtype=np.int16)
            return array

    @staticmethod
    def hasTasksInfo(idx):
        """查询任务对象"""
        try:
            task = AllinpayRobotQiniuLog.objects.get(id=idx, status=0)
        except BaseException as err:
            return None
        else:
            return task

    @staticmethod
    def hasAudioAttr(idx):
        """查询声音模型"""
        try:
            audio = AllinpayRobotConfig.objects.get(id=idx, status=1)
        except BaseException as err:
            return None
        else:
            return audio

    @staticmethod
    def doFormatString(text):
        """格式化字符串"""
        out = (
            text.strip()
            .replace(r"\n\n", "。")
            .replace(r"\n", "。")
            .replace("*", "")
            .replace("#", "")
            .replace("/", "")
            .replace(" ", "")
            .replace("-", "")
            .replace('"', "")
            .replace("'", "")
            .replace("“", "")
            .replace("”", "")
            .replace("’", "")
            .replace("‘", "")
            .replace(":", "。")
            .replace("：", "。")
            .replace("?", "。")
            .replace("？", "。")
            .replace("!", "。")
            .replace("！", "。")
            .replace("。。", "。")
        )
        return out
